import cv2
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import numpy as np
from PIL import Image
from config import FIRE_DETECTION_CONFIDENCE, FIRE_DETECTION_MODEL_PATH

class FireDetectionModel(nn.Module):
    def __init__(self):
        super(FireDetectionModel, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        self.classifier = nn.Sequential(
            nn.Linear(256 * 28 * 28, 512),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(512, 2)
        )

    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

class FireDetector:
    def __init__(self):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model = FireDetectionModel().to(self.device)
        self.load_model()
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
        ])

    def load_model(self):
        """加载预训练的火灾检测模型"""
        try:
            self.model.load_state_dict(torch.load(FIRE_DETECTION_MODEL_PATH,
                                                map_location=self.device))
            self.model.eval()
        except Exception as e:
            print(f"无法加载火灾检测模型: {str(e)}")
            print("使用未训练的模型继续运行")

    def detect(self, frame):
        """检测图像中是否存在火灾场景"""
        try:
            # 将OpenCV格式转换为PIL格式
            image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            image = Image.fromarray(image)
            
            # 预处理图像
            input_tensor = self.transform(image)
            input_batch = input_tensor.unsqueeze(0).to(self.device)

            # 进行预测
            with torch.no_grad():
                output = self.model(input_batch)
                probabilities = torch.softmax(output, dim=1)
                fire_prob = probabilities[0][1].item()

            # 如果火灾概率超过阈值，在图像上标记警告
            if fire_prob >= FIRE_DETECTION_CONFIDENCE:
                cv2.putText(frame, f'Fire Detected ({fire_prob:.2f})',
                            (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1,
                            (0, 0, 255), 2)
                return True

            return False

        except Exception as e:
            print(f"火灾检测过程中发生错误: {str(e)}")
            return False